# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional
from fastapi import HTTPException
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.docstore.document import Document
from datetime import datetime, timedelta
import os
import shutil
import threading
from ..config import settings

class ChatService:
    """聊天服务类，封装所有与聊天相关的核心功能"""
    
    def __init__(self):
        """初始化聊天服务"""
        self.sessions: Dict[str, Dict[str, Any]] = {}  # 会话存储
        self.embedding = OpenAIEmbeddings()
        self.lock = threading.RLock()  # 可重入锁，提高并发安全性
        self._start_session_cleanup()  # 启动会话清理线程
    
    def _start_session_cleanup(self) -> None:
        """启动会话清理线程，定期清理超时会话"""
        def cleanup_sessions():
            """清理超时会话的函数"""
            while True:
                try:
                    with self.lock:
                        now = datetime.now()
                        timeout = timedelta(hours=settings.SESSION_TIMEOUT_HOURS)
                        expired_sessions = [
                            session_id for session_id, session in self.sessions.items()
                            if now - session["last_activity"] > timeout
                        ]
                        
                        # 清理超时会话
                        for session_id in expired_sessions:
                            self._delete_session_data(session_id)
                            del self.sessions[session_id]
                            print(f"清理超时会话: {session_id}")
                
                except Exception as e:
                    print(f"会话清理出错: {str(e)}")
                
                # 每隔1小时检查一次
                threading.Event().wait(3600)
        
        # 启动后台线程
        cleanup_thread = threading.Thread(target=cleanup_sessions, daemon=True)
        cleanup_thread.start()
    
    def _delete_session_data(self, session_id: str) -> None:
        """删除会话相关数据"""
        # 删除向量数据库文件
        chroma_path = os.path.join(settings.CHROMA_PERSIST_DIR, session_id)
        if os.path.exists(chroma_path):
            shutil.rmtree(chroma_path, ignore_errors=True)
        
        # 删除文档文件
        docs_path = os.path.join(settings.DOCUMENTS_DIR, session_id)
        if os.path.exists(docs_path):
            shutil.rmtree(docs_path, ignore_errors=True)
    
    def create_session(self, 
                      agent_name: str, 
                      agent_description: str, 
                      model_name: str = settings.DEFAULT_MODEL, 
                      temperature: float = settings.DEFAULT_TEMPERATURE) -> str:
        """创建新会话"""
        with self.lock:
            # 生成会话ID
            session_id = f"sess_{int(datetime.now().timestamp() * 1000)}_{os.urandom(4).hex()}"
            
            # 初始化会话目录
            os.makedirs(os.path.join(settings.DOCUMENTS_DIR, session_id), exist_ok=True)
            
            # 存储会话信息
            self.sessions[session_id] = {
                "vectorstore": None,
                "memory": ConversationBufferMemory(
                    memory_key="chat_history", 
                    return_messages=True
                ),
                "documents": [],
                "agent_name": agent_name,
                "agent_description": agent_description,
                "model_name": model_name,
                "temperature": temperature,
                "last_activity": datetime.now()
            }
            
            return session_id
    
    def get_vectorstore(self, session_id: str) -> Chroma:
        """获取或创建会话的向量数据库"""
        with self.lock:
            if session_id not in self.sessions:
                raise HTTPException(status_code=404, detail="会话不存在或已过期")
            
            # 更新活动时间
            self.sessions[session_id]["last_activity"] = datetime.now()
            
            # 如果向量数据库不存在，创建一个新的
            if self.sessions[session_id]["vectorstore"] is None:
                chroma_path = os.path.join(settings.CHROMA_PERSIST_DIR, session_id)
                self.sessions[session_id]["vectorstore"] = Chroma(
                    embedding_function=self.embedding,
                    persist_directory=chroma_path
                )
            
            return self.sessions[session_id]["vectorstore"]
    
    def get_chat_chain(self, session_id: str) -> ConversationalRetrievalChain:
        """获取会话的聊天链"""
        session = self._get_session(session_id)
        
        # 创建聊天模型
        llm = ChatOpenAI(
            model_name=session["model_name"],
            temperature=session["temperature"],
            streaming=True,
            max_tokens=settings.DEFAULT_MAX_TOKENS
        )
        
        # 创建检索链
        retriever = self.get_vectorstore(session_id).as_retriever()
        
        # 创建对话检索链
        return ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever,
            memory=session["memory"],
            return_source_documents=True
        )
    
    def _get_session(self, session_id: str) -> Dict[str, Any]:
        """获取会话信息，不存在则抛出异常"""
        with self.lock:
            if session_id not in self.sessions:
                raise HTTPException(status_code=404, detail="会话不存在或已过期")
            
            # 更新活动时间
            self.sessions[session_id]["last_activity"] = datetime.now()
            return self.sessions[session_id]
    
    def process_document(self, file_path: str, file_type: str) -> List[Document]:
        """处理不同类型的文档并返回文档内容"""
        try:
            if file_type == "application/pdf":
                from langchain.document_loaders import PyPDFLoader
                loader = PyPDFLoader(file_path)
            elif file_type == "text/plain":
                from langchain.document_loaders import TextLoader
                loader = TextLoader(file_path)
            elif "word" in file_type:
                from langchain.document_loaders import UnstructuredWordDocumentLoader
                loader = UnstructuredWordDocumentLoader(file_path)
            elif "powerpoint" in file_type:
                from langchain.document_loaders import UnstructuredPowerPointLoader
                loader = UnstructuredPowerPointLoader(file_path)
            else:
                raise ValueError(f"不支持的文件类型: {file_type}")
            
            documents = loader.load()
            
            # 分割文档为片段
            from langchain.text_splitter import RecursiveCharacterTextSplitter
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200,
                length_function=len
            )
            return text_splitter.split_documents(documents)
        
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}")
    
    def upload_document(self, 
                       session_id: str, 
                       file_path: str, 
                       file_name: str, 
                       file_type: str) -> Dict[str, Any]:
        """上传文档并添加到向量数据库"""
        try:
            # 处理文档
            documents = self.process_document(file_path, file_type)
            if not documents:
                return {
                    "status": "warning",
                    "message": f"文档 '{file_name}' 处理成功，但未提取到内容",
                    "document_count": 0
                }
            
            # 获取向量数据库
            vectorstore = self.get_vectorstore(session_id)
            
            # 添加文档到向量数据库
            with self.lock:
                # 如果已有文档，合并
                if self.sessions[session_id]["documents"]:
                    existing_documents = self.sessions[session_id]["documents"]
                    all_documents = existing_documents + documents
                    self.sessions[session_id]["documents"] = all_documents
                    
                    # 重新创建向量数据库
                    chroma_path = os.path.join(settings.CHROMA_PERSIST_DIR, session_id)
                    if os.path.exists(chroma_path):
                        shutil.rmtree(chroma_path, ignore_errors=True)
                    
                    vectorstore = Chroma.from_documents(
                        documents=all_documents,
                        embedding=self.embedding,
                        persist_directory=chroma_path
                    )
                    self.sessions[session_id]["vectorstore"] = vectorstore
                else:
                    # 新建向量数据库
                    vectorstore.add_documents(documents)
                    self.sessions[session_id]["documents"] = documents
            
            return {
                "status": "success",
                "message": f"文档 '{file_name}' 已成功上传并处理",
                "document_count": len(documents)
            }
        
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"上传失败: {str(e)}")
    
    def chat(self, session_id: str, message: str) -> Dict[str, Any]:
        """处理聊天消息并返回响应"""
        # 获取聊天链
        chat_chain = self.get_chat_chain(session_id)
        
        # 执行聊天链
        try:
            return chat_chain({"question": message})
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"聊天处理失败: {str(e)}")
    
    def clear_history(self, session_id: str) -> Dict[str, str]:
        """清除会话历史"""
        session = self._get_session(session_id)
        
        with self.lock:
            # 清除记忆
            session["memory"].clear()
            
            return {
                "status": "success",
                "message": "会话历史已清除"
            }
    
    def delete_session(self, session_id: str) -> Dict[str, str]:
        """删除会话"""
        with self.lock:
            if session_id not in self.sessions:
                raise HTTPException(status_code=404, detail="会话不存在或已过期")
            
            # 删除会话数据
            self._delete_session_data(session_id)
            
            # 从会话中移除
            del self.sessions[session_id]
            
            return {
                "status": "success",
                "message": "会话已删除"
            }
    
    def get_session_info(self, session_id: str) -> Dict[str, Any]:
        """获取会话信息"""
        session = self._get_session(session_id)
        
        return {
            "status": "success",
            "session_info": {
                "agent_name": session["agent_name"],
                "agent_description": session["agent_description"],
                "model_name": session["model_name"],
                "temperature": session["temperature"],
                "document_count": len(session.get("documents", [])),
                "last_activity": session["last_activity"].isoformat()
            }
        }    